# dataset consts
ORIGIN_DATA_PATH = '../origin_data'
DATA_PATH = '../data'
DATA_CONF = 'conf.json'
LOG_PATH = '../logs'
DOC_PATH = 'docs'

# lark const
LARK_HEADERS = {'Content-Type': 'application/json;charset=utf-8'}
LARK_API_DIR = 'conf/lark_api.json'
LARK_INTERVL = 60*10  # 10min

# task conf consts
NUM_CLASS = 10
TASK_CONF_PATH = './conf'
INI_SECTION = 'basic'
BASIC_CONF_DICT = {
    'dataset': 'cifar_10',
    'distr': 'ALL',
    'lr': '2e-4',
    'batch': '72',
    'epoch': '40',
    'task': 'stageA'
}
OPT_DISTR = ['ALL', 'LTL', 'MLT']
OPT_TASKS = {
    'stageG': {},
    'stageA': {
        'backbone': 'resnet101'
    },
    'stageB': {
        'backbone': 'resnet101',
        'model': None,
        'gan': None,
        'method': "TRANS,CBLoss,FocalLoss,SMOTE,GAN"
    }
}


def get_distr(distr_type: str, mask: list[int] = [1]*NUM_CLASS):
    distr = [x for x in mask]
    if distr_type == 'ALL':
        return distr
    elif distr_type == 'LTL':
        i = 1
        for j in range(len(distr)):
            if distr[j] != 0:
                distr[j] = 1/i
                i += 1
        return distr
    elif distr_type == 'MTL':
        num_cls = sum(distr)
        mean = sum([1/(i+1) for i in range(num_cls)])/num_cls
        for i in range(len(distr)):
            if distr[i] != 0:
                distr[i] = mean
        return distr


def split(line: str):
    return line.split(',')


CONF_TYPE = {
    'dataset': str,
    'distr': get_distr,
    'lr': float,
    'batch': int,
    'epoch': int,
    'task': str,
    'backbone': str,
    'method': split,
    'model': str,
    'gan': str
}
